import tool
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import mapping
import mydraw

plt.rcParams["font.sans-serif"] = ["SimHei"]  # 使用黑体
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题


def partial_corr_analysis():
    data = []
    for i in range(1, 7):  # 假设有6个产品
        product_data = tool.get_np(f"type_monthly_total/{i}.csv")["monthly_total"]
        data.append(product_data)

    # Step 1: 计算六类商品之间的相关系数矩阵
    data = np.array(data).T  # 转置以使每一列代表一个商品
    R = np.corrcoef(data, rowvar=False)  # 相关系数矩阵

    # Step 2: 计算商品与总销量之间的相关系数
    total_sales = np.sum(data, axis=1)  # 总销量
    r_XiT = []
    for i in range(data.shape[1]):
        r = np.corrcoef(data[:, i], total_sales)[0, 1]  # 商品与总销量的相关系数
        r_XiT.append(r)
    r_XiT = np.array(r_XiT)

    # Step 3: 计算偏相关系数
    partial_corr_matrix = np.zeros((6, 6))
    for i in range(6):
        for j in range(6):
            if i != j:
                partial_corr_matrix[i, j] = (R[i, j] - r_XiT[i] * r_XiT[j]) / np.sqrt(
                    (1 - r_XiT[i] ** 2) * (1 - r_XiT[j] ** 2)
                )
            else:
                partial_corr_matrix[i, j] = 1.0  # 对角线上的偏相关系数为1

    print("偏相关系数矩阵：")
    print(partial_corr_matrix)
    # 类别名称
    labels = mapping.type_map1.values()
    # 创建热力图
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        partial_corr_matrix,
        annot=True,
        cmap="coolwarm",
        xticklabels=labels,
        yticklabels=labels,
        vmin=-1,
        vmax=1,
        cbar_kws={"label": ""},
    )
    plt.yticks(rotation=0)
    # 添加标题
    plt.title("偏相关系数矩阵（控制总销量）")
    # 显示图表
    mydraw.show_or_print("partical_corr.png")


def partical_corr_product():
    data = []
    pro_list = [5, 6, 7, 10, 15, 16]
    for i in pro_list:  # 假设有6个产品
        product_data = tool.get_np(f"product_monthly_total/{i}.csv")["monthly_total"]
        data.append(product_data)

    data = np.array(data).T  # 转置以使每一列代表一个商品
    R = np.corrcoef(data, rowvar=False)  # 相关系数矩阵

    # Step 2: 计算商品与总销量之间的相关系数
    total_sales = np.sum(data, axis=1)  # 总销量
    r_XiT = []
    for i in range(data.shape[1]):
        r = np.corrcoef(data[:, i], total_sales)[0, 1]  # 商品与总销量的相关系数
        r_XiT.append(r)
    r_XiT = np.array(r_XiT)

    # Step 3: 计算偏相关系数
    partial_corr_matrix = np.zeros((6, 6))
    for i in range(6):
        for j in range(6):
            if i != j:
                partial_corr_matrix[i, j] = (R[i, j] - r_XiT[i] * r_XiT[j]) / np.sqrt(
                    (1 - r_XiT[i] ** 2) * (1 - r_XiT[j] ** 2)
                )
            else:
                partial_corr_matrix[i, j] = 1.0  # 对角线上的偏相关系数为1

    # print("偏相关系数矩阵：")
    print(partial_corr_matrix)
    # 类别名称
    labels = [mapping.get_name(i) for i in pro_list]
    # 创建热力图
    plt.figure(figsize=(8, 6))
    sns.heatmap(
        partial_corr_matrix,
        annot=True,
        cmap="coolwarm",
        xticklabels=labels,
        yticklabels=labels,
        vmin=-1,
        vmax=1,
        cbar_kws={"label": ""},
    )
    plt.yticks(rotation=0)
    # 添加标题
    plt.title("偏相关系数矩阵(产品)")
    # plt.show()
    print(123)
    # 显示图表
    mydraw.show_or_print("partical_corr_product.png")
    # mydraw.show_or_print()


if __name__ == "__main__":
    partical_corr_product()
